/te-ci pytorch

Greptile Summary

This PR improves the device-init grouped linear module by adding full support for a single grouped weight/bias parameter (single_grouped_weight / single_grouped_bias) through the native GroupedTensor GEMM backend. It also introduces an explicit grouped_gemm_backend argument to GroupedLinear, deprecates the NVTE_GROUPED_LINEAR_USE_FUSED_GROUPED_GEMM environment variable, and fixes an MXFP8 scale-address bug in the group_quantize_mxfp8 CUDA kernel for single-tensor uniform-group layouts.

• Single grouped weight/bias: _GroupedLinear._prepare_weights_for_grouped_tensor_gemm and _prepare_bias_for_grouped_tensor_gemm handle both discrete per-expert tensors and a single stacked GroupedTensor parameter; the backward mirrors this by recovering per-group gradients via rowwise_data views.
• MXFP8 weight caching: tex.group_quantize gains an output= reuse argument (with quantizer-identity and shape validation in C++) and a noop_flag path so CUDA-graph-captured weight quantization can be replayed in-place without pointer changes.
• Bug fix – out_features = weights[0].size(-2) replaces .size(0), which incorrectly returned the expert count for a 3-D grouped-weight tensor.

Confidence Score: 4/5

The changes are well-structured and covered by targeted tests; the main risks are the unguarded quantizer-internal flag mutation and test helpers that would silently skip on future GPU generations.

The core forward/backward plumbing for single grouped parameters is logically consistent and the new tests exercise the key paths (BF16, MXFP8, fuse_wgrad, delay_wgrad, return_bias). The weight_quantizer.internal flag is mutated without a try/finally guard — an exception during quantization would leave it in the wrong state for all subsequent calls. The device-capability upper bound <= (11, 0) in both test helpers is tighter than the production gating, meaning tests would silently skip on SM >= 12.x hardware. The CUDA kernel integer division first_logical_dim / num_tensors lacks an assertion, though the invariant is enforced upstream.

transformer_engine/pytorch/module/grouped_linear.py (quantizer-internal exception safety) and the two test helper functions in test_grouped_linear.py and test_grouped_mlp.py (device capability upper bound).

Important Files Changed

Sequence Diagram

sequenceDiagram
    participant Caller
    participant GroupedLinear_Module
    participant _GroupedLinear_Autograd
    participant PrepareWeights
    participant tex_group_quantize
    participant GroupedGEMM

    Caller->>GroupedLinear_Module: forward(x, m_splits, is_first_microbatch)
    GroupedLinear_Module->>GroupedLinear_Module: resolve grouped_gemm_backend
    GroupedLinear_Module->>GroupedLinear_Module: get weight/bias tensors (single or discrete)
    GroupedLinear_Module->>_GroupedLinear_Autograd: "apply(inp, m_splits, non_tensor_args, *weights, *biases)"

    _GroupedLinear_Autograd->>PrepareWeights: "_prepare_weights_for_grouped_tensor_gemm(single_grouped_weight=True)"
    alt BF16 weight
        PrepareWeights-->>_GroupedLinear_Autograd: GroupedTensorStorage (view)
    else MXFP8 weight, workspace cached
        PrepareWeights->>tex_group_quantize: "group_quantize(src, quantizer, N, noop_flag, output=workspace)"
        tex_group_quantize-->>PrepareWeights: workspace (in-place, same pointers)
        PrepareWeights-->>_GroupedLinear_Autograd: cached GroupedTensorStorage
    else MXFP8 weight, no workspace
        PrepareWeights->>tex_group_quantize: group_quantize(src, quantizer, N)
        tex_group_quantize-->>PrepareWeights: new GroupedTensorStorage
        PrepareWeights-->>_GroupedLinear_Autograd: new workspace (saved to _fp8_workspaces)
    end

    _GroupedLinear_Autograd->>GroupedGEMM: general_grouped_gemm_for_grouped_tensor(weight, x, out)
    GroupedGEMM-->>_GroupedLinear_Autograd: output tensor

    _GroupedLinear_Autograd-->>GroupedLinear_Module: (output, new_workspaces)
    GroupedLinear_Module->>GroupedLinear_Module: update _fp8_workspaces["weight"]

    alt return_bias and single_grouped_bias
        GroupedLinear_Module-->>Caller: (output, bias_GroupedTensor)
    else
        GroupedLinear_Module-->>Caller: output
    end

_{Reviews (1): Last reviewed commit: "lint" | Re-trigger Greptile}